
from .utils import convert_arch


class Shellcode:
    """
    Base class of all shellcode snippets.

    :ivar list os:      All OSes that this piece of shellcode applies to.
    :ivar archses:      All architectures that this piece of shellcode runs on. Those architecture strings should be
                        parsable by archinfo.arch_from_id().
    :type archses:      list[str]
    :ivar str name:     Name of the shellcode.
    :ivar str asm:      Assembly code of this piece of shellcode.
    :ivar bytes code:   Byte-representation of the shellcode. `asm` will be used if `code` is None.
    """

    os = [ ]
    arches = [ ]
    name = None
    asm = None
    code = None

    def __init__(self, *args, **kwargs):
        # do nothing. just make sure Shellcode() can be initialized with arbitrary arguments without crashing.
        pass

    def __repr__(self):
        return "<Shellcode %s on %s[%s]>" % (self.name, "/".join(self.os), "/".join(self.arches))

    @staticmethod
    def check_shellcode_for_incompatible_chars(shellcode_bytes):
        bs = shellcode_bytes
        dangerous_chars_found = set()
        if any(c in bs for c in [0x00, 0x0a, 0x20, 0x25, 0x2b, 0x2d, 0x3b]):
            for i, c in enumerate(bs):
                if c in [0x00, 0x0a, 0x20, 0x25, 0x2b, 0x2d, 0x3b]:
                    dangerous_chars_found.add(c)
                    print(f"{hex(i)}: {bytes([c])}[{hex(c)}]")
            print('This shellcode is dangerous to use in HTTP payloads')
        return dangerous_chars_found

    def raw(self, arch=None):
        """
        return a raw string representing the shellcode, subclasses which more sophisticated
        shellcodes will want to take options here.

        :param arch:        Architecture (or its name). Unnecessary if `self.code` is provided or the shellcode only has
                            one supported architecture.
        :type arch:         str or archinfo.Arch or None
        :return:            Byte representation of the shellcode snippet.
        :rtype:             bytes
        """

        if self.code:
            return self.code

        assert self.asm is not None

        # assemble shellcode from the given instructions
        if len(self.arches) == 1:
            the_arch = self.arches[0]
        elif arch is not None:
            the_arch = convert_arch(arch)
        else:
            raise ValueError("Please specify the architecture where the shellcode snippet will apply on.")
        code = the_arch.asm(self.code)
        return code
